import os
import requests
import json
import gradio as gr
import uuid
import logging
import time  # 新增时间模块
from datetime import datetime
from dotenv import load_dotenv
from typing import List, Dict, Optional, Generator
from functools import lru_cache
from pydantic import BaseModel, ValidationError, Field

# ==================== 日志配置 ====================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("WenxinAssistant")

# ==================== 数据模型 ====================
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatSession(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    title: str = Field(default="新对话")
    create_time: str = Field(default_factory=lambda: datetime.now().isoformat())
    messages: List[ChatMessage] = []
    

# ==================== 初始化配置 ====================
load_dotenv()
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
HISTORY_FILE = os.path.join(BASE_DIR, "chat_sessions.json")
MAX_RETRIES = 3
API_TIMEOUT = 300

class ConfigError(Exception):
    """自定义配置异常"""
    pass

# ==================== 文心一言封装类（优化版）====================
class ErnieLLM:
    def __init__(self):
        self.API_KEY = os.getenv("BAIDU_WENXIN_APP_KEY")
        self.SECRET_KEY = os.getenv("BAIDU_WENXIN_SECRET_KEY")
        
        if not all([self.API_KEY, self.SECRET_KEY]):
            raise ConfigError("缺少必要的环境变量配置，请检查.env文件")
            
        self.token_expire_time = 0
        self.expires_in = 2592000  # 默认有效期
        self.ACCESS_TOKEN = None
        self._refresh_token()

    def _get_access_token_with_retry(self) -> str:
        """动态获取有效期并处理错误"""
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials",
            "client_id": self.API_KEY,
            "client_secret": self.SECRET_KEY,
        }
        
        for attempt in range(MAX_RETRIES):
            try:
                response = requests.post(url, params=params, timeout=API_TIMEOUT)
                response.raise_for_status()
                token_data = response.json()
                self.expires_in = token_data["expires_in"]  # 动态获取有效期
                return token_data["access_token"]
            except KeyError as e:
                logger.error(f"Token响应格式错误: {str(e)}")
                raise
            except requests.exceptions.RequestException as e:
                logger.error(f"获取Token失败（尝试 {attempt+1}/{MAX_RETRIES}）: {str(e)}")
                if attempt == MAX_RETRIES - 1:
                    raise RuntimeError("无法获取API访问令牌") from e

    def _refresh_token(self):
        """带缓存清理的Token刷新"""
        logger.info("正在刷新Access Token...")
        try:
            self.ACCESS_TOKEN = self._get_access_token_with_retry()
            self.token_expire_time = time.time() + self.expires_in - 60
            self._cached_chat.cache_clear()  # 清除所有缓存
            logger.info(f"Token刷新成功，有效期至 {time.strftime('%Y-%m-%d %H:%M', time.localtime(self.token_expire_time))}")
        except Exception as e:
            logger.critical("Token刷新失败，服务不可用")
            raise

    def _check_token_valid(self):
        """带重试的Token检查"""
        if time.time() > self.token_expire_time:
            logger.warning("检测到Token过期，正在自动刷新...")
            self._refresh_token()

    @lru_cache(maxsize=100)
    def _cached_chat(self, message: str, history_hash: int) -> Optional[str]:
        """实际缓存实现"""
        return None  # 此处应为真实API调用，示例简化处理

    def chat(self, message: str, history: List[Dict], temperature: float = 0.7, max_tokens: int = 500) -> str:
        """全功能聊天方法"""
        try:
            self._check_token_valid()
            
            # 生成缓存键
            cache_key = hash(json.dumps({
                "message": message,
                "history": history,
                "temp": temperature,
                "tokens": max_tokens
            }, sort_keys=True))
            
            # 缓存检查
            if cached_response := self._cached_chat(message, cache_key):
                logger.info("命中缓存响应")
                return cached_response

            # API请求
            url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token={self.ACCESS_TOKEN}"
            headers = {"Content-Type": "application/json"}
            payload = json.dumps({
                "messages": history + [{"role": "user", "content": message}],
                "temperature": temperature,
                "max_tokens": max_tokens
            })

            for attempt in range(MAX_RETRIES):
                try:
                    response = requests.post(url, headers=headers, data=payload, timeout=API_TIMEOUT)
                    response.raise_for_status()
                    result = response.json()["result"]
                    
                    # 更新缓存（示例需实现真实缓存逻辑）
                    return result
                
                except requests.exceptions.HTTPError as e:
                    error_code = response.json().get("error_code", 0)
                    if error_code == 110:  # 百度API的Token错误码
                        logger.warning("Token已失效，触发自动刷新")
                        self._refresh_token()
                        return self.chat(message, history, temperature, max_tokens)
                    else:
                        raise
                except requests.exceptions.Timeout:
                    logger.error(f"请求超时（尝试 {attempt+1}/{MAX_RETRIES}）")
                    if attempt == MAX_RETRIES - 1:
                        return "⚠️ 请求超时，请稍后重试"
                except Exception as e:
                    logger.error(f"请求失败（尝试 {attempt+1}/{MAX_RETRIES}）: {str(e)}")
                    if attempt == MAX_RETRIES - 1:
                        return f"⚠️ 请求失败：{str(e)}"
        except Exception as e:
            logger.error(f"系统异常: {str(e)}")
            return f"⚠️ 系统错误：{str(e)}"

# ==================== 应用核心逻辑（优化版）==================== 
ernie_bot = ErnieLLM() #ErnieLLM封装了会话大模型


def validate_and_fix_session(session: Dict) -> Optional[ChatSession]:
    """严格的数据校验和修复"""
    try:
        # 处理旧版数据格式，类型转换逻辑
        if "messages" in session:
            # 列表类型校验：确保messages是列表且不为空时进行检查
            if isinstance(session["messages"], list) and len(session["messages"]) > 0:
                # # 字典转模型对象：转换旧字典格式消息为ChatMessage对象
                if isinstance(session["messages"][0], dict):
                    session["messages"] = [
                        ChatMessage(**msg) for msg in session["messages"]
                        if "role" in msg and "content" in msg
                    ]
            else:
                # 处理空消息列表的情况
                session["messages"] = []
                
        return ChatSession(**session)
    except ValidationError as e:
        logger.warning(f"无效的会话数据: {str(e)}")
        return None

def save_sessions(sessions: List[ChatSession]) -> None:
    """带数据校验的保存方法（原子操作）"""
    try:
        # 使用临时文件避免写入中断
        temp_file = f"{HISTORY_FILE}.tmp"
        
        valid_sessions = [s.model_dump() for s in sessions if isinstance(s, ChatSession)]
        
        with open(temp_file, "w", encoding="utf-8") as f:
            json.dump(valid_sessions, f, 
                     ensure_ascii=False, 
                     indent=2,
                     default=str)
        
        # 原子替换文件
        if os.path.exists(temp_file):
            os.replace(temp_file, HISTORY_FILE)
            
    except Exception as e:
        logger.error(f"保存会话失败: {str(e)}")
        if os.path.exists(temp_file):
            os.remove(temp_file)
        raise

def load_sessions() -> List[ChatSession]:
    """增强版历史加载"""
    sessions = []
    if os.path.exists(HISTORY_FILE):
        try:
            with open(HISTORY_FILE, "r", encoding="utf-8") as f:
                raw_data = json.load(f)
                for item in raw_data:
                    # 添加空数据过滤
                    if not isinstance(item, dict):
                        continue
                        
                    # 确保包含必要字段
                    item.setdefault("id", str(uuid.uuid4()))
                    item.setdefault("title", "未命名对话")
                    item.setdefault("create_time", datetime.now().isoformat())
                    item.setdefault("messages", [])
                    
                    if valid := validate_and_fix_session(item):
                        sessions.append(valid)
        except (json.JSONDecodeError, KeyError) as e:
            logger.error(f"历史文件损坏: {str(e)}")
    return sessions

def create_session() -> ChatSession:
    """创建带完整验证的新会话"""
    return ChatSession(
        id=str(uuid.uuid4()),
        title=datetime.now().strftime("%m月%d日 %H:%M 对话"),
        create_time=datetime.now().isoformat(),
        messages=[]
    )

# ==================== Gradio界面（优化版）====================
with gr.Blocks(
    title="文心一言智能助手-优化版",
    theme=gr.themes.Soft(primary_hue="blue"),
    css="""
    .message {
        padding: 12px; 
        margin: 8px 0; 
        border-radius: 15px;
        max-width: 100%;
        transition: transform 0.2s;
    }
    .user-message {
        background: #f0f7ff; 
        margin-left: auto;
        border: 1px solid #cce0ff;
    }
    .assistant-message {
        background: #f8f9fa;
        border: 1px solid #e0e0e0;
    }
    .textbox-input {
        font-size: 14px;
        padding: 12px;
    }
    .error-message {
        color: #dc3545;
        border-color: #dc3545 !important;
    }
    """
) as app:
    current_session = gr.State()
    all_sessions = gr.State(load_sessions)
    
    gr.Markdown("""<div style="text-align:center"><h1>🌸 文心智能助手 专业版</h1></div>""")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                height=600,
                avatar_images=(
                    ("user.png" if os.path.exists("user.png") else None),
                    ("ai.png" if os.path.exists("ai.png") else None)
                ),
                elem_classes=["message-container"],
                type="messages"  # 修改点2：添加type参数
            )
            msg = gr.Textbox(
                placeholder="请输入消息（Enter换行，Shift+Enter发送）...", 
                lines=2,
                max_lines=5,
                elem_classes="textbox-input"
            )
            with gr.Row():
                submit_btn = gr.Button("🚀 发送", variant="primary")
                new_btn = gr.Button("🆕 新建对话")
                clear_btn = gr.Button("🧹 清空当前")

        with gr.Column(scale=1):
            with gr.Accordion("⚙️ 参数设置", open=False):
                temperature = gr.Slider(0.1, 1.0, value=0.7, label="回答随机性")
                max_length = gr.Slider(100, 1000, value=500, label="最大长度")
            
            with gr.Accordion("📜 历史对话", open=True):
                session_dropdown = gr.Dropdown(
                    label="对话列表",
                    allow_custom_value=True,
                    interactive=True
                )
                with gr.Row():
                    refresh_btn = gr.Button("🔄 刷新")
                    delete_btn = gr.Button("❌ 删除", variant="stop")

    # ==================== 事件处理（优化版）====================
    def get_dropdown_choices(sessions: List[ChatSession]):
        return [
            (f"{datetime.fromisoformat(s.create_time).strftime('%m/%d %H:%M')} | {s.title}", s.id)
            for s in sessions
        ]

    def init_state():
        sessions = load_sessions()
        # 强制创建初始会话
        if not sessions or any(not isinstance(s, ChatSession) for s in sessions):
            new_session = create_session()
            sessions = [new_session]
            save_sessions(sessions)
        return sessions[0], sessions, gr.update(choices=get_dropdown_choices(sessions))

    def new_session_action(sessions: List[ChatSession]):
        new_session = create_session()
        sessions.insert(0, new_session)
        save_sessions(sessions)
        return new_session, sessions, gr.update(choices=get_dropdown_choices(sessions))

    def delete_session_action(sessions: List[ChatSession], selected_id: str):
        if selected_id:
            sessions = [s for s in sessions if s.id != selected_id]
            # 确保总有一个有效会话
            if not sessions:
                new_session = create_session()
                sessions = [new_session]
            current = sessions[0]
            save_sessions(sessions)
            return current, sessions, gr.update(choices=get_dropdown_choices(sessions))
        return sessions[0] if sessions else create_session(), sessions, gr.skip()

    # 修改点4：加强空值检查
    def update_chat_messages(session: Optional[ChatSession]):
        if session and session.messages:
            return [{"role": msg.role, "content": msg.content} for msg in session.messages]
        return []

    def process_message(message: str, session: Optional[ChatSession]) -> ChatSession:
        """消息预处理"""
        if session is None:
            session = create_session()
            
        if not message.strip():
            raise gr.Error("消息内容不能为空")
            
        session.messages.append(ChatMessage(role="user", content=message.strip()))
        return session

    def generate_response_optimized(
        message: str,
        session: Optional[ChatSession],
        sessions: List[ChatSession],
        temp: float,
        max_len: int
    ) -> Generator:
        try:
             # 新增会话有效性检查
            if session is None:
                session = create_session()
                sessions = [session]
                save_sessions(sessions)

            # 处理用户输入
            session = process_message(message, session)
            yield update_chat_messages(session), session, sessions, ""
            
            # 添加等待状态
            session.messages.append(ChatMessage(role="assistant", content="思考中..."))
            yield update_chat_messages(session), session, sessions, ""
            
            # 生成响应
            response = ernie_bot.chat(
                message=message,
                history=[msg.model_dump() for msg in session.messages[:-1]],  # 修改点5：使用model_dump()
                temperature=temp,
                max_tokens=max_len
            )
            
            # 更新会话记录
            session.messages[-1] = ChatMessage(role="assistant", content=response)
            
            # 限制历史长度
            if len(session.messages) > 100:
                session.messages = session.messages[-50:]
            
            # 更新会话列表
            updated_sessions = [
                s if s.id != session.id else session
                for s in sessions
            ]
            
            save_sessions(updated_sessions)
            yield update_chat_messages(session), session, updated_sessions, ""
            
        except Exception as e:
            logger.error(f"生成响应时出错: {str(e)}")
            error_msg = ChatMessage(role="assistant", content=f"⚠️ 系统错误：{str(e)}")
            session.messages.append(error_msg)
            yield update_chat_messages(session), session, sessions, ""

    # ==================== 事件绑定 ====================
    app.load(
        init_state,
        outputs=[current_session, all_sessions, session_dropdown]
    )
    
    submit_triggers = [
        msg.submit(
            generate_response_optimized,
            inputs=[msg, current_session, all_sessions, temperature, max_length],
            outputs=[chatbot, current_session, all_sessions, msg]
        ),
        submit_btn.click(
            generate_response_optimized,
            inputs=[msg, current_session, all_sessions, temperature, max_length],
            outputs=[chatbot, current_session, all_sessions, msg]
        )
    ]

    new_btn.click(
        new_session_action,
        [all_sessions],
        [current_session, all_sessions, session_dropdown]
    )
    
    delete_btn.click(
        delete_session_action,
        [all_sessions, session_dropdown],
        [current_session, all_sessions, session_dropdown]
    )
    
    refresh_btn.click(
        lambda: gr.update(choices=get_dropdown_choices(load_sessions())),
        outputs=session_dropdown
    )
    
    clear_btn.click(
        lambda s: (create_session(), []),
        [current_session],
        [current_session, chatbot]
    )
    
    session_dropdown.input(
        lambda sessions, selected_id: next((s for s in sessions if s.id == selected_id), None),
        [all_sessions, session_dropdown],
        [current_session]
    ).then(
        update_chat_messages,
        [current_session],
        [chatbot]
    )

if __name__ == "__main__":
    try:
        if not os.path.exists(HISTORY_FILE):
            with open(HISTORY_FILE, "w") as f:
                json.dump([create_session().model_dump()], f)
        app.launch(server_port=7860)
    except Exception as e:
        logger.critical(f"程序启动失败: {str(e)}")
        raise